import pymysql
from datetime import datetime

# 数据库配置
db_config = {
    'host': '192.168.88.66',
    'user': 'root',
    'password': '202011',
    'charset': 'utf8mb4',
    'cursorclass': pymysql.cursors.DictCursor
}


# 创建数据库和表
def initialize_database():
    connection = pymysql.connect(**db_config)

    try:
        with connection.cursor() as cursor:
            # 创建数据库
            cursor.execute("CREATE DATABASE IF NOT EXISTS ctf_platform")
            cursor.execute("USE ctf_platform")

            # 创建表
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS instances (
                id INT AUTO_INCREMENT PRIMARY KEY,
                container_id VARCHAR(64) NOT NULL,
                target_name VARCHAR(50) NOT NULL,
                ports TEXT NOT NULL,
                start_time DATETIME NOT NULL,
                end_time DATETIME,
                status ENUM('running', 'stopped') DEFAULT 'running',
                completed BOOLEAN DEFAULT 0,
                user_ip VARCHAR(45),
                UNIQUE KEY (container_id)
            )""")

            # 添加测试数据
            cursor.execute("""
            INSERT INTO instances (container_id, target_name, ports, start_time, status, completed)
            VALUES 
                ('test123', 'web1', '{\"80/tcp\": 6001}', %s, 'stopped', 1),
                ('test456', 'pwn1', '{\"9999/tcp\": 6002}', %s, 'running', 0)
            """, (datetime.now(), datetime.now()))

        connection.commit()
        print("Database initialized successfully!")

    finally:
        connection.close()


if __name__ == '__main__':
    initialize_database()